package com.luncheng.mcp.core.router;

import com.luncheng.mcp.core.McpSseServerBoostrap;
import io.modelcontextprotocol.server.McpSyncServer;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.RouterFunctions;
import org.springframework.web.servlet.function.ServerResponse;

import java.util.Map;

@Slf4j
@Configuration
//@RefreshScope
public class DynamicRouterProvider {

    private final MutableRouterFunction mutableRouterFunction = new MutableRouterFunction();

    @PostConstruct
    public void init() {
        log.info("Initializing dynamic router...");
        refreshRouter();
    }

    public void refreshRouter() {
        try {
            RouterFunction<ServerResponse> newRouter = RouterFunctions.route()
                    .path("", builder -> {
                        Map<String, McpSyncServer> serverMap = McpSseServerBoostrap.getServerMap();
                        if (serverMap == null || serverMap.isEmpty()) {
                            log.warn("No MCP servers found when building dynamic routes.");
                            builder.GET("/health", req -> ServerResponse.ok().body("OK"));
                            return;
                        }
                        serverMap.forEach((name, server) -> {
                            var provider = McpSseServerBoostrap.getProviderMap().get(name);
                            if (provider != null) {
                                log.info("Registering dynamic route for MCP server: {}", name);
                                builder.add(provider.getRouterFunction());
                            } else {
                                log.warn("No provider found for MCP server: {}", name);
                            }
                        });
                    })
                    .build();

            mutableRouterFunction.setDelegate(newRouter);
            log.info("Router function has been refreshed dynamically.");
        } catch (Exception e) {
            log.error("Failed to build dynamic router function", e);
        }
    }

    @Bean
    public RouterFunction<ServerResponse> dynamicRouterFunction() {
        return this.mutableRouterFunction;
    }
}